'''
Created on Aug 10, 2009

@author: mkiyer
'''

from veggie.genome.interval import GenomicInterval
from veggie.genome.io import GenomicIntervalReader
import sys

class ParseError(Exception):
    def __init__(self, *args, **kwargs):
        Exception.__init__(self, *args)
        self.linenum = kwargs.get("linenum", None)
    def __str__(self):
        if self.linenum:
            return Exception.__str__(self) + " on line " + str(self.linenum)
        else:
            return Exception.__str__(self)

class GenomicIntervalReader(object):
    """
    Reader for iterating a set of intervals in a tab separated file. Can
    also parse header and comment lines if requested.
    
    Code based on bx-python's GenomicIntervalReader    
    """

    def __init__(self, input, 
                 chrom_col=0, start_col=1, end_col=2, strand_col=5, 
                 default_strand='+',
                 comment_lines_startswith = ["#", "track"]):
        """
        Defaults to read BED format
        """
        self.input = input
        self.input_iter = iter(input)
        self.chrom_col = chrom_col
        self.start_col = start_col
        self.end_col = end_col
        self.strand_col = strand_col        
        self.default_strand = default_strand
        self.comment_lines_startswith = comment_lines_startswith
        self.linenum = 0
        self.print_debug = True
        self.print_line_count = 0
        self.print_every_n_lines = 100000
        
    def __iter__(self):
        return self
      
    def next(self):
        thisline = self.input_iter.next()
        self.linenum += 1
        # print every n lines for debugging
        if self.print_debug:
            self.print_line_count += 1
            if self.print_line_count == self.print_every_n_lines:
                sys.stderr.write(" %d\n" % (self.linenum))
        # check to make sure line can be parsed            
        if not thisline:
            raise ParseError("line is None")
        thisline = thisline.strip()
        # catch blank lines
        if not thisline:
            # TODO: warning?
            raise ParseError("blank line")
        # check for comments
        for comment_line_start in self.comment_lines_startswith:
            if thisline.startswith(comment_line_start):
                raise ParseError("comment line")
        try:
            for interval in self.parse_line(thisline):
                yield interval
        except ParseError, e:
            e.linenum = self.linenum
            raise e
    
    def parse_line(self, thisline):
        fields = thisline.strip.split('\t')
        return [GenomicInterval(chrom=fields[self.chrom_col],
                                chrom_start=fields[self.start_col],
                                chrom_end=fields[self.end_col],
                                strand=fields[self.strand_col])]

class BowtieReader(GenomicIntervalReader):
    '''
    The bowtie aligner outputs each alignment on a separate line. Each line 
    is a collection of 8 fields separated by tabs; from left to right, 
    the fields are:

    1. Name of read that aligned
    2. Orientation of read in the alignment, - for reverse complement, 
        + otherwise
    3. Name of reference sequence where alignment occurs, or ordinal ID if 
        no name was provided
    4. 1-based offset into the forward reference strand where leftmost 
        character of the alignment occurs
    5. Read sequence (reverse-complemented if orientation is -)
    6. ASCII-encoded read qualities (reversed if orientation is -). 
       The encoded quality values are on the Phred scale and the 
       encoding is ASCII-offset by 33 (ASCII char !).
    7. Number of other instances where the same read aligns against the same
       reference characters as were aligned against in this alignment. This 
       is not the number of other places the read aligns with the same 
       number of mismatches. The number in this column is generally not a 
       good proxy for that number (e.g., the number in this column may be 
       '0' while the number of other alignments with the same number of 
       mismatches might be large). This column was previously described as 
       "Reserved".
    8. Comma-separated list of mismatch descriptors. If there are no 
       mismatches in the alignment, this field is empty. A single descriptor 
       has the format offset:reference-base>read-base. The offset is 
       expressed as a 0-based offset from the high-quality (5') end of 
       the read.
    '''
    def __init__(self, input):
        super(BowtieReader, self).__init__(input, comment_lines_startswith=[])        
        self.prev_interval = None
        self.end_of_input = False

    def __iter__(self):
        while True:
            yield self.gen_group()
    
    def gen_group(self):
        while True:
            # get a line of input
            thisline = self.input_iter.next()

            #try:
            #    thisline = self.input_iter.next()
            #except StopIteration:
            #    if self.prev_interval != None:
            #        yield self.prev_interval
            #    return
                
            # parent class has a check_line function to check for
            # errors and keep track of line number, etc
            # TODO: checkline
            # thisline = self.check_line(thisline)
            try:
                interval = self.parse_line(thisline)[0]
            except ParseError, e:
                e.linenum = self.linenum
                raise e

            if self.prev_interval == None:
                self.prev_interval = interval
            else:
                yield self.prev_interval
                tmp_prev_interval = self.prev_interval
                self.prev_interval = interval
                if interval.name != tmp_prev_interval.name:
                    return    
                
    def parse_line(self, thisline):        
        # split the line into fields and extract relevant information
        fields = thisline.split('\t')
        name = fields[0]
        strand = fields[1]
        # convert the chromosome mapping field to ucsc chromosome format
        chrom = fields[2].split('|')[-1]
        if chrom.startswith('NC'):
            chrom = str(int(chrom.split('_')[-1]))                    
            if chrom == '23':
                chrom = 'X'
            elif chrom == '24':
                chrom = 'Y'
            chrom = 'chr' + str(chrom)
        chrom_start = int(fields[3]) - 1
        chrom_end = chrom_start + len(fields[4])
        # TODO: handle mismatches
        if chrom:
            interval = GenomicInterval(chrom=chrom,
                                       chrom_start=chrom_start,
                                       chrom_end=chrom_end,
                                       strand=strand)
            # add name attribute for grouping mult-mapping reads
            interval.name = name
            return [interval]
        else:
            raise ParseError(linenum=self.linenum)


if __name__ == '__main__':
    i = [1, 1, 1, 2, 2, 5, 6, 7, 8, 8, 8, 8, 9]    
    
    for group in A(i):
        for v in group:
            print v    
